package proxy

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

type Server struct {
	engine *gin.Engine
	closeCh chan struct{}
	upgrader *websocket.Upgrader
	running int32
}

func NewServer () *Server {
	return &Server{
		engine: gin.Default(),
		closeCh: make(chan struct{}, 1),
		upgrader: &websocket.Upgrader{
			HandshakeTimeout: 10 * time.Second,
		},
	}
}

func (s *Server) Start(port int) {
	if s.running == 1 {
		return
	}

	s.engine.NoRoute(func(c *gin.Context) {
		forward(c, s)
	})
	gin.SetMode(gin.DebugMode)
	go s.engine.Run(fmt.Sprintf(":%d", port))

	atomic.StoreInt32(&s.running, int32(1))
}

func forward(c *gin.Context, s *Server) {
	if (websocket.IsWebSocketUpgrade(c.Request)) {
		webSocketForward(c, s)
	} else {
		httpForward(c, s)	
	}
}

func httpForward(c *gin.Context, s *Server) {
	request := c.Request
	writer := c.Writer
	backendServer := BackendServers.getRandomServer()

	httpProxy := NewHttpProxy(backendServer)
	forwardResponse, err := httpProxy.HttpForward(request)
	if err != nil {
		fmt.Printf("HttpForward error, %v\n", err)
		c.AbortWithError(500, err)
		return
	}

	writer.WriteHeader(forwardResponse.StatusCode)
	for key, value := range forwardResponse.Header {
		writer.Header().Add(key, value[0])
	}

	data, err := io.ReadAll(forwardResponse.Body)
	if err != nil {
		fmt.Printf("get forwardResponse error, %v", err)
		writer.Write([]byte(err.Error()))
	} else {
		writer.Write(data)
		defer forwardResponse.Body.Close()
	}

	writer.Flush()
}

func webSocketForward(c *gin.Context, s *Server) {
	request := c.Request
	writer := c.Writer
	backendServer := BackendServers.getRandomServer()

	outboundConn, err := s.upgrader.Upgrade(writer, request, writer.Header())
	if err != nil {
		fmt.Printf("websocket handshake error, %v", err)
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}
	webSocketProxy, err := NewWebSocketProxy(outboundConn, backendServer, request)
	if err != nil {
		fmt.Printf("NewWebSocketProxy error, %v", err)
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}		

	go webSocketProxy.Start()

	LOOP:
	for {
		messageType, data, err := outboundConn.ReadMessage()
		if err != nil {
			log.Printf("outboundConn readmessage error, err=%v", err)
			webSocketProxy.Stop()
			break
		}

		switch messageType {
		case websocket.CloseMessage:
			webSocketProxy.Stop()
			outboundConn.Close()
			break LOOP
					
		case websocket.PingMessage:
			webSocketProxy.SendToProxy(messageType, data)
		case websocket.TextMessage:
			webSocketProxy.SendToProxy(messageType, data)
		}	
	}
}

func (s *Server) Stop() {
	close(s.closeCh)
}